import os
#from http://code.activestate.com/recipes/286222/

_proc_status = '/proc/%d/status' % os.getpid()

_scale = {'kB': 1024.0, 'mB': 1024.0*1024.0,
          'KB': 1024.0, 'MB': 1024.0*1024.0}

def _memory():
    global _proc_status, _scale
     # get pseudo file  /proc/<pid>/status
    try:
        t = open(_proc_status)
        s = t.read()
        t.close()
    except:
        return 0.0  # non-Linux?
    # get VmKey line e.g. 'VmRSS:  9999  kB\n ...'
    r = []
    for key in ['VmSize:', 'VmRSS:', 'VmStk:']:
        i = s.index(key)
        v = s[i:].split(None, 3)  # whitespace
        if len(v) < 3:
            r.append(0.0)
        else:
            r.append(int (float(v[1]) * _scale[v[2]]) )
    return r



"""fnl.core.memory module

Trivial, but working code to get the memory usage of the current process
where the pid is retrieved using os.getpid() and the memory usage is read
from the unix command ps.    

http://snipplr.com/view.php?codeview&id=6460
"""

#__version__ = "1.0"
#__author__ = "Florian Leitner"

def _mem(size="rss"):
    """Generalization; memory sizes: rss, rsz, vsz."""
    return int(os.popen('ps -p %d -o %s | tail -1' %
                        (os.getpid(), size)).read())


def memory():
    try:
        return _memory()
    except Exception:
        pass

    try:
        # FIXME: third number currently doesnt match in both methods (size (text, data, stack) vs stack size)
        res = []
        res.append(int(_mem('vsz')*1024.0))
        res.append(int(_mem('rsz')*1024.0))
        res.append(int(_mem('sz')*1024.0))
        return res 

    except Exception:
        pass
    return 0,0,0

#if __name__ == '__main__':
#    print _memory()
#    print [int(_mem('vsz')*1024.0), int(_mem('rsz')*1024.0), int(_mem('sz')*1024.0)]

